from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import scipy.io as sio
import numpy as np
from scipy.signal import hilbert, welch
import io
import csv
import os

app = Flask(__name__)
CORS(app)

# 配置
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def calculate_rms(signal):
    """计算RMS值"""
    return np.sqrt(np.mean(np.square(signal)))

def calculate_peak(signal):
    """计算峰值"""
    return np.max(np.abs(signal))

def calculate_kurtosis(signal):
    """计算峭度"""
    mean = np.mean(signal)
    std = np.std(signal)
    return np.mean(((signal - mean) / std) ** 4)

def calculate_spectrum(signal, fs):
    """计算频谱"""
    f, psd = welch(signal, fs, nperseg=min(1024, len(signal)//4))
    return f, psd

@app.route('/api/upload', methods=['POST'])
def upload_file():
    """上传并处理MAT文件"""
    if 'file' not in request.files:
        return jsonify({'error': '没有文件被上传'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if file and file.filename.endswith('.mat'):
        filename = file.filename
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            # 读取MAT文件
            mat_data = sio.loadmat(filepath)
            
            # 假设数据在'X'变量中，如果不存在则寻找第一个数值变量
            signal_key = None
            for key in mat_data.keys():
                if not key.startswith('__') and isinstance(mat_data[key], np.ndarray):
                    signal_key = key
                    break
            
            if signal_key is None:
                return jsonify({'error': '未找到有效的信号数据'}), 400
            
            signal = mat_data[signal_key].flatten()
            fs = 12000  # 采样频率12kHz
            
            # 计算特征
            rms = calculate_rms(signal)
            peak = calculate_peak(signal)
            kurtosis = calculate_kurtosis(signal)
            
            # 计算频谱
            f, psd = calculate_spectrum(signal, fs)
            
            # 计算包络谱
            analytic_signal = hilbert(signal)
            envelope = np.abs(analytic_signal)
            f_env, psd_env = calculate_spectrum(envelope, fs)
            
            # 准备返回数据
            time = np.arange(len(signal)) / fs
            
            response = {
                'filename': filename,
                'time': time.tolist(),
                'signal': signal.tolist(),
                'frequency': f.tolist(),
                'psd': psd.tolist(),
                'envelope_frequency': f_env.tolist(),
                'envelope_psd': psd_env.tolist(),
                'features': {
                    'rms': float(rms),
                    'peak': float(peak),
                    'kurtosis': float(kurtosis),
                    'max_frequency': float(f[np.argmax(psd)]),
                    'max_amplitude': float(np.max(psd))
                }
            }
            
            return jsonify(response)
            
        except Exception as e:
            return jsonify({'error': f'处理文件时出错: {str(e)}'}), 500
    
    return jsonify({'error': '请上传.mat格式的文件'}), 400

@app.route('/api/export_csv', methods=['POST'])
def export_csv():
    """导出CSV报告"""
    data = request.json
    
    # 创建CSV文件
    output = io.StringIO()
    writer = csv.writer(output)
    
    # 写入标题
    writer.writerow(['参数', '值'])
    writer.writerow(['文件名', data.get('filename', '')])
    writer.writerow(['RMS', data.get('rms', '')])
    writer.writerow(['峰值', data.get('peak', '')])
    writer.writerow(['峭度', data.get('kurtosis', '')])
    writer.writerow(['主频', data.get('max_frequency', '')])
    writer.writerow(['主频幅值', data.get('max_amplitude', '')])
    writer.writerow([])
    
    # 写入时域数据
    writer.writerow(['时间(s)', '幅值'])
    time_data = data.get('time', [])
    signal_data = data.get('signal', [])
    for t, s in zip(time_data, signal_data):
        writer.writerow([t, s])
    
    # 准备文件下载
    output.seek(0)
    return send_file(
        io.BytesIO(output.getvalue().encode()),
        mimetype='text/csv',
        as_attachment=True,
        download_name='fault_analysis_report.csv'
    )

if __name__ == '__main__':
    app.run(debug=True, port=5000)
